import pymysql
import pandas as pd
from typing import Optional

def get_connection(host='localhost', port=3306, user='root', password='', database=''):
    """
    创建并返回MySQL数据库连接
    
    Args:
        host: 数据库主机地址
        port: 数据库端口
        user: 用户名
        password: 密码
        database: 数据库名
    
    Returns:
        pymysql.connections.Connection: 数据库连接对象
    """
    connection = pymysql.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
        charset='utf8mb4'
    )
    return connection


def create_fund_data_table(connection):
    """
    创建基金数据表
    
    Args:
        connection: MySQL数据库连接对象
        
    Returns:
        bool: 创建成功返回True，失败返回False
    """
    try:
        cursor = connection.cursor()
        
        # 创建表的SQL语句
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS fund_data (
            id BIGINT AUTO_INCREMENT PRIMARY KEY COMMENT '主键ID',
            fund_code VARCHAR(20) NOT NULL COMMENT '基金代码',
            fund_name VARCHAR(100) COMMENT '基金简称',
            nav_date DATE NOT NULL COMMENT '净值日期',
            unit_nav DECIMAL(10, 4) COMMENT '单位净值',
            accumulated_nav DECIMAL(10, 4) COMMENT '累计净值',
            nav_abnormal TINYINT(1) COMMENT '净值是否异常(0:正常, 1:异常)',
            
            -- 1周持有期指标
            return_1w DECIMAL(10, 2) COMMENT '持有1周收益率(%)',
            max_drawdown_1w DECIMAL(10, 2) COMMENT '持有1周最大回撤(%)',
            profit_prob_1w DECIMAL(10, 2) COMMENT '持有1周盈利概率',
            return_mean_1w DECIMAL(10, 2) COMMENT '持有1周收益率均值',
            return_std_1w DECIMAL(10, 2) COMMENT '持有1周收益率标准差',
            drawdown_mean_1w DECIMAL(10, 2) COMMENT '持有1周最大回撤均值',
            drawdown_std_1w DECIMAL(10, 2) COMMENT '持有1周最大回撤标准差',
            
            -- 1月持有期指标
            return_1m DECIMAL(10, 2) COMMENT '持有1月收益率(%)',
            max_drawdown_1m DECIMAL(10, 2) COMMENT '持有1月最大回撤(%)',
            profit_prob_1m DECIMAL(10, 2) COMMENT '持有1月盈利概率',
            return_mean_1m DECIMAL(10, 2) COMMENT '持有1月收益率均值',
            return_std_1m DECIMAL(10, 2) COMMENT '持有1月收益率标准差',
            drawdown_mean_1m DECIMAL(10, 2) COMMENT '持有1月最大回撤均值',
            drawdown_std_1m DECIMAL(10, 2) COMMENT '持有1月最大回撤标准差',
            
            -- 3月持有期指标
            return_3m DECIMAL(10, 2) COMMENT '持有3月收益率(%)',
            max_drawdown_3m DECIMAL(10, 2) COMMENT '持有3月最大回撤(%)',
            profit_prob_3m DECIMAL(10, 2) COMMENT '持有3月盈利概率',
            return_mean_3m DECIMAL(10, 2) COMMENT '持有3月收益率均值',
            return_std_3m DECIMAL(10, 2) COMMENT '持有3月收益率标准差',
            drawdown_mean_3m DECIMAL(10, 2) COMMENT '持有3月最大回撤均值',
            drawdown_std_3m DECIMAL(10, 2) COMMENT '持有3月最大回撤标准差',
            
            -- 半年持有期指标
            return_6m DECIMAL(10, 2) COMMENT '持有半年收益率(%)',
            max_drawdown_6m DECIMAL(10, 2) COMMENT '持有半年最大回撤(%)',
            profit_prob_6m DECIMAL(10, 2) COMMENT '持有半年盈利概率',
            return_mean_6m DECIMAL(10, 2) COMMENT '持有半年收益率均值',
            return_std_6m DECIMAL(10, 2) COMMENT '持有半年收益率标准差',
            drawdown_mean_6m DECIMAL(10, 2) COMMENT '持有半年最大回撤均值',
            drawdown_std_6m DECIMAL(10, 2) COMMENT '持有半年最大回撤标准差',
            
            -- 1年持有期指标
            return_1y DECIMAL(10, 2) COMMENT '持有1年收益率(%)',
            max_drawdown_1y DECIMAL(10, 2) COMMENT '持有1年最大回撤(%)',
            profit_prob_1y DECIMAL(10, 2) COMMENT '持有1年盈利概率',
            return_mean_1y DECIMAL(10, 2) COMMENT '持有1年收益率均值',
            return_std_1y DECIMAL(10, 2) COMMENT '持有1年收益率标准差',
            drawdown_mean_1y DECIMAL(10, 2) COMMENT '持有1年最大回撤均值',
            drawdown_std_1y DECIMAL(10, 2) COMMENT '持有1年最大回撤标准差',
            
            -- 2年持有期指标
            return_2y DECIMAL(10, 2) COMMENT '持有2年收益率(%)',
            max_drawdown_2y DECIMAL(10, 2) COMMENT '持有2年最大回撤(%)',
            profit_prob_2y DECIMAL(10, 2) COMMENT '持有2年盈利概率',
            return_mean_2y DECIMAL(10, 2) COMMENT '持有2年收益率均值',
            return_std_2y DECIMAL(10, 2) COMMENT '持有2年收益率标准差',
            drawdown_mean_2y DECIMAL(10, 2) COMMENT '持有2年最大回撤均值',
            drawdown_std_2y DECIMAL(10, 2) COMMENT '持有2年最大回撤标准差',
            
            -- 3年持有期指标
            return_3y DECIMAL(10, 2) COMMENT '持有3年收益率(%)',
            max_drawdown_3y DECIMAL(10, 2) COMMENT '持有3年最大回撤(%)',
            profit_prob_3y DECIMAL(10, 2) COMMENT '持有3年盈利概率',
            return_mean_3y DECIMAL(10, 2) COMMENT '持有3年收益率均值',
            return_std_3y DECIMAL(10, 2) COMMENT '持有3年收益率标准差',
            drawdown_mean_3y DECIMAL(10, 2) COMMENT '持有3年最大回撤均值',
            drawdown_std_3y DECIMAL(10, 2) COMMENT '持有3年最大回撤标准差',
            
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
            
            INDEX idx_fund_code (fund_code),
            INDEX idx_nav_date (nav_date),
            UNIQUE KEY uk_fund_date (fund_code, nav_date)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='基金数据表';
        """
        
        cursor.execute(create_table_sql)
        connection.commit()
        print("✓ 基金数据表创建成功或已存在")
        cursor.close()
        return True
        
    except Exception as e:
        print(f"✗ 创建表失败: {e}")
        return False


def insert_fund_data(connection, df: pd.DataFrame) -> bool:
    """
    将基金数据插入到MySQL数据库
    
    Args:
        connection: MySQL数据库连接对象
        df: 包含基金数据的DataFrame
        
    Returns:
        bool: 插入成功返回True，失败返回False
    """
    if df is None or df.empty:
        print("⚠ 警告: DataFrame为空，无数据可插入")
        return False
    
    try:
        cursor = connection.cursor()
        
        # 定义列名映射关系（DataFrame列名 -> 数据库字段名）
        column_mapping = {
            '基金代码': 'fund_code',
            '基金简称': 'fund_name',
            '净值日期': 'nav_date',
            '单位净值': 'unit_nav',
            '累计净值': 'accumulated_nav',
            '净值异常': 'nav_abnormal',
            
            '持有1周收益率(%)': 'return_1w',
            '持有1周最大回撤(%)': 'max_drawdown_1w',
            '持有1周盈利概率': 'profit_prob_1w',
            '持有1周收益率均值': 'return_mean_1w',
            '持有1周收益率标准差': 'return_std_1w',
            '持有1周最大回撤均值': 'drawdown_mean_1w',
            '持有1周最大回撤标准差': 'drawdown_std_1w',
            
            '持有1月收益率(%)': 'return_1m',
            '持有1月最大回撤(%)': 'max_drawdown_1m',
            '持有1月盈利概率': 'profit_prob_1m',
            '持有1月收益率均值': 'return_mean_1m',
            '持有1月收益率标准差': 'return_std_1m',
            '持有1月最大回撤均值': 'drawdown_mean_1m',
            '持有1月最大回撤标准差': 'drawdown_std_1m',
            
            '持有3月收益率(%)': 'return_3m',
            '持有3月最大回撤(%)': 'max_drawdown_3m',
            '持有3月盈利概率': 'profit_prob_3m',
            '持有3月收益率均值': 'return_mean_3m',
            '持有3月收益率标准差': 'return_std_3m',
            '持有3月最大回撤均值': 'drawdown_mean_3m',
            '持有3月最大回撤标准差': 'drawdown_std_3m',
            
            '持有半年收益率(%)': 'return_6m',
            '持有半年最大回撤(%)': 'max_drawdown_6m',
            '持有半年盈利概率': 'profit_prob_6m',
            '持有半年收益率均值': 'return_mean_6m',
            '持有半年收益率标准差': 'return_std_6m',
            '持有半年最大回撤均值': 'drawdown_mean_6m',
            '持有半年最大回撤标准差': 'drawdown_std_6m',
            
            '持有1年收益率(%)': 'return_1y',
            '持有1年最大回撤(%)': 'max_drawdown_1y',
            '持有1年盈利概率': 'profit_prob_1y',
            '持有1年收益率均值': 'return_mean_1y',
            '持有1年收益率标准差': 'return_std_1y',
            '持有1年最大回撤均值': 'drawdown_mean_1y',
            '持有1年最大回撤标准差': 'drawdown_std_1y',
            
            '持有2年收益率(%)': 'return_2y',
            '持有2年最大回撤(%)': 'max_drawdown_2y',
            '持有2年盈利概率': 'profit_prob_2y',
            '持有2年收益率均值': 'return_mean_2y',
            '持有2年收益率标准差': 'return_std_2y',
            '持有2年最大回撤均值': 'drawdown_mean_2y',
            '持有2年最大回撤标准差': 'drawdown_std_2y',
            
            '持有3年收益率(%)': 'return_3y',
            '持有3年最大回撤(%)': 'max_drawdown_3y',
            '持有3年盈利概率': 'profit_prob_3y',
            '持有3年收益率均值': 'return_mean_3y',
            '持有3年收益率标准差': 'return_std_3y',
            '持有3年最大回撤均值': 'drawdown_mean_3y',
            '持有3年最大回撤标准差': 'drawdown_std_3y',
        }
        
        # 获取DataFrame中存在的列
        available_columns = [col for col in column_mapping.keys() if col in df.columns]
        db_columns = [column_mapping[col] for col in available_columns]
        
        # 构建INSERT语句
        placeholders = ', '.join(['%s'] * len(db_columns))
        columns_str = ', '.join(db_columns)
        
        # 使用ON DUPLICATE KEY UPDATE实现插入或更新
        update_str = ', '.join([f"{col}=VALUES({col})" for col in db_columns if col not in ['fund_code', 'nav_date']])
        
        insert_sql = f"""
        INSERT INTO fund_data ({columns_str})
        VALUES ({placeholders})
        ON DUPLICATE KEY UPDATE {update_str}
        """
        
        # 批量插入数据
        success_count = 0
        fail_count = 0
        
        for idx, row in df.iterrows():
            try:
                # 提取数据并处理NaN值
                values = []
                for col in available_columns:
                    value = row[col]
                    # 处理pandas的NaN和NaT
                    if pd.isna(value):
                        values.append(None)
                    elif col == '净值异常':
                        # 布尔值转换为0或1
                        values.append(1 if value else 0)
                    else:
                        values.append(value)
                
                cursor.execute(insert_sql, values)
                success_count += 1
                
            except Exception as e:
                fail_count += 1
                print(f"✗ 插入第 {idx} 行数据失败: {e}")
        
        connection.commit()
        cursor.close()
        
        print(f"✓ 数据插入完成: 成功 {success_count} 条，失败 {fail_count} 条")
        return True
        
    except Exception as e:
        print(f"✗ 插入数据失败: {e}")
        connection.rollback()
        return False


def save_fund_data_to_db(df: pd.DataFrame, 
                         host='localhost', 
                         port=3306, 
                         user='root', 
                         password='Aa123456.', 
                         database='fund_data') -> bool:
    """
    将基金数据保存到MySQL数据库（一站式函数）
    
    Args:
        df: 包含基金数据的DataFrame
        host: 数据库主机地址
        port: 数据库端口
        user: 用户名
        password: 密码
        database: 数据库名
        
    Returns:
        bool: 保存成功返回True，失败返回False
    """
    connection = None
    try:
        # 创建数据库连接
        connection = get_connection(host=host, port=port, user=user, password=password, database=database)
        
        # 创建表（如果不存在）
        if not create_fund_data_table(connection):
            return False
        
        # 插入数据
        result = insert_fund_data(connection, df)
        
        return result
        
    except Exception as e:
        print(f"✗ 保存数据到数据库失败: {e}")
        return False
        
    finally:
        if connection:
            connection.close()


# 使用示例
if __name__ == "__main__":
    # 创建连接
    conn = get_connection(
        host='localhost',               #数据库主机地址
        user='root',                    #用户名
        password='Aa123456.',           #密码
        database='fund_data'            #数据库名
    )
    
    # 使用连接
    cursor = conn.cursor()
    
    # 查看数据库中有哪些表
    cursor.execute("SHOW TABLES")
    tables = cursor.fetchall()
    print("数据库中的表：")
    for table in tables:
        print(f"  - {table[0]}")

    
    # 关闭连接
    cursor.close()
    conn.close()
